import numpy as np
import scipy.spatial.distance
import matplotlib.pyplot as plt
from sklearn.cluster import KMeans
from mykmeans import MyKMeans
from utils2 import JS_D
import time

centre_path = 'E:/Research/2020ContrastiveLearningForSceneLabel/Data/20210329ExperimentData/round2/ds_cluster/centre2.npy'
centre = np.load(centre_path)
fig = plt.figure(figsize=(25,12))
ax = fig.add_subplot(1, 3, 1)
plt.bar(range(100), centre[0],width=1, align='edge')
plt.ylim(0,0.08)
ax.set_xticks([0,25,50,75,99])
ax.set_xticklabels(['-1','-0.5','0','0.5','1'])

ax = fig.add_subplot(1, 3, 2)
plt.bar(range(100), centre[1],width=1, align='edge')
plt.ylim(0,0.08)
ax.set_xticks([0,25,50,75,99])
ax.set_xticklabels(['-1','-0.5','0','0.5','1'])

ax = fig.add_subplot(1, 3, 3)
plt.bar(range(100), centre[2],width=1, align='edge')
plt.ylim(0,0.08)
ax.set_xticks([0,25,50,75,99])
ax.set_xticklabels(['-1','-0.5','0','0.5','1'])


fig.suptitle(centre_path)
plt.savefig('E:/Research/2020ContrastiveLearningForSceneLabel/Data/20210329ExperimentData/round2/ds_cluster/centre2.png')
plt.show()


exit(0) #######################################################################################################


ds_path = 'E:/Research/2020ContrastiveLearningForSceneLabel/Data/20210329ExperimentData/round2/round2_2ds.npy'
ds_feat = np.load(ds_path)
for i in range(50):
    print(i)
    start = time.time()
    n_clusters = 3
    estimator = KMeans(n_clusters=n_clusters)
    estimator.fit(ds_feat)
    label_pred = estimator.labels_
    centre = np.zeros((n_clusters, 100))
    for j in range(n_clusters):
        centre[j] = np.average(ds_feat[label_pred == j], axis=0)

    print('K-means time: ', time.time() - start)

    np.save('E:/Research/2020ContrastiveLearningForSceneLabel/Data/20210329ExperimentData/round2/ds_cluster_eucliden/centre%d.npy'%i, centre)

    fig = plt.figure(figsize=(25,12))
    ax = fig.add_subplot(1, 3, 1)
    plt.bar(range(100), centre[0],width=1, align='edge')
    plt.ylim(0,0.08)
    ax.set_xticks([0,25,50,75,99])
    ax.set_xticklabels(['-1','-0.5','0','0.5','1'])

    ax = fig.add_subplot(1, 3, 2)
    plt.bar(range(100), centre[1],width=1, align='edge')
    plt.ylim(0,0.08)
    ax.set_xticks([0,25,50,75,99])
    ax.set_xticklabels(['-1','-0.5','0','0.5','1'])

    ax = fig.add_subplot(1, 3, 3)
    plt.bar(range(100), centre[2],width=1, align='edge')
    plt.ylim(0,0.08)
    ax.set_xticks([0,25,50,75,99])
    ax.set_xticklabels(['-1','-0.5','0','0.5','1'])

    plt.savefig('E:/Research/2020ContrastiveLearningForSceneLabel/Data/20210329ExperimentData/round2/ds_cluster_eucliden/centre%d.png'%i)
    # plt.show()
    plt.close()

exit(0) #######################################################################################################



ds_path = 'E:/Research/2020ContrastiveLearningForSceneLabel/Data/20210329ExperimentData/round2/round2_2ds.npy'
ds_feat = np.load(ds_path)
for i in range(50):
    print(i)
    start = time.time()
    estimator = MyKMeans(n_clusters=3, distance_metric= JS_D)
    estimator.fit(ds_feat)
    label_pred = estimator.labels_
    centre = estimator.centre
    print('K-means time: ', time.time() - start)

    np.save('E:/Research/2020ContrastiveLearningForSceneLabel/Data/20210329ExperimentData/round2/ds_cluster/centre%d.npy'%i, centre)

    fig = plt.figure(figsize=(25,12))
    ax = fig.add_subplot(1, 3, 1)
    plt.bar(range(100), centre[0],width=1, align='edge')
    plt.ylim(0,0.08)
    ax.set_xticks([0,25,50,75,99])
    ax.set_xticklabels(['-1','-0.5','0','0.5','1'])

    ax = fig.add_subplot(1, 3, 2)
    plt.bar(range(100), centre[1],width=1, align='edge')
    plt.ylim(0,0.08)
    ax.set_xticks([0,25,50,75,99])
    ax.set_xticklabels(['-1','-0.5','0','0.5','1'])

    ax = fig.add_subplot(1, 3, 3)
    plt.bar(range(100), centre[2],width=1, align='edge')
    plt.ylim(0,0.08)
    ax.set_xticks([0,25,50,75,99])
    ax.set_xticklabels(['-1','-0.5','0','0.5','1'])

    plt.savefig('E:/Research/2020ContrastiveLearningForSceneLabel/Data/20210329ExperimentData/round2/ds_cluster/centre%d.png'%i)
    # plt.show()
    plt.close()

exit(0) #######################################################################################################